import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft


def generate_sinusoid(N, A, f0, fs, phi):
    """
    N(int) : number of samples
    A(float) : amplitude
    f0(float): frequency in Hz
    fs(float): sample rate
    phi(float): initial phase
    :return: X(numpy array): sinusoid signal which length is M
    """
    T = 1 / fs
    n = np.arange(N)
    x = A * np.cos(2 * f0 * np.pi * n * T + phi)
    return x
# 另一种生成正弦信号的方法，生成时长为t的序列
def generate_sinusoid_2(t, A, f0, fs, phi):
    """
    t  (float) : 生成序列的时长
    A  (float) : amplitude
    f0 (float) : frequency
    fs (float) : sample rate
    phi(float) : initial phase
    returns
    x (numpy array): sinusoid signal sequence
    """
    T = 1.0 / fs
    N = t / T
    return generate_sinusoid(N, A, f0, fs, phi)



if __name__ == '__main__':
    # generate_sinusoid
    # N = 511
    # A = 0.8
    # f0 = 440
    # fs = 44100
    # phi = 0
    # x = generate_sinusoid(N, A, f0, fs, phi)
    # plt.plot(x)
    # plt.show()


    # A = 1.0
    # f0 = 440
    # fs = 44100
    # phi = 0
    # t = 0.02
    #
    # x = generate_sinusoid_2(t, A, f0, fs, phi)
    #
    # n = np.arange(0, 0.02, 1 / fs)
    # plt.plot(n, x)
    # plt.show()
    # generate sinusoid
    N = 511
    A = 0.8
    f0 = 440
    fs = 44100
    phi = 1.0
    x = generate_sinusoid(N, A, f0, fs, phi)

    # fft
    X = fft(x)
    mX = np.abs(X)  # magnitude
    pX = np.angle(X) # phase

    # plot the magnitude and phase
    plt.subplot(2, 1, 1)
    plt.plot(mX)

    plt.subplot(2, 1, 2)
    plt.plot(pX)
    plt.show()

